//
//  icmp.cpp
//  DeviceLock
//
//  Created by shb on 3/4/13.
//
//

#include "icmp.h"

#include <boost/asio.hpp>
#include <boost/bind.hpp>

#include "icmp_header.hpp"
#include "ipv4_header.hpp"

using boost::asio::ip::icmp;
using boost::asio::deadline_timer;
namespace posix_time = boost::posix_time;

class pinger
{
public:
    pinger(boost::asio::io_service& io_service, const char* destination, unsigned int attempts, IPingObserver* observer)
    : resolver_(io_service), socket_(io_service, icmp::v4()),
    timer_(io_service), sequence_number_(0), num_replies_(0),
    attempts_(attempts), observer_(observer), expires_at_(0)
    {
        assert(attempts_);
        assert(observer_);
        
        icmp::resolver::query query(icmp::v4(), destination, "");
        destination_ = *resolver_.resolve(query);
        
        start_send();
        start_receive();
    }
    
    // packet expires time is needed to determine when ping timeout exceeds
    unsigned int expires_at()
    {
        static unsigned int s_expires = expires_at_ * milliseconds_multiplier_;
        return s_expires;
    }
    
private:
    void start_send()
    {
        if (attempts_-- == 0)
        {
            resolver_.get_io_service().stop();
            observer_->OnComplete();
            return;
        }
        
        
        std::string body("DL Service ping.");
        
        // Create an ICMP header for an echo request.
        icmp_header echo_request;
        echo_request.type(icmp_header::echo_request);
        echo_request.code(0);
        echo_request.identifier(get_identifier());
        echo_request.sequence_number(++sequence_number_);
        compute_checksum(echo_request, body.begin(), body.end());
        
        // Encode the request packet.
        boost::asio::streambuf request_buffer;
        std::ostream os(&request_buffer);
        os << echo_request << body;
        
        // Send the request.
        time_sent_ = posix_time::microsec_clock::universal_time();
        socket_.send_to(request_buffer.data(), destination_);
        
        // Wait up to five seconds (expires_at) for a reply.
        num_replies_ = 0;
        expires_at_ = 5;
        timer_.expires_at(time_sent_ + posix_time::seconds(expires_at_));
        timer_.async_wait(boost::bind(&pinger::handle_timeout, this));
    }
    
    void handle_timeout()
    {
        if (num_replies_ == 0)
        {
            posix_time::ptime now = posix_time::microsec_clock::universal_time();
            observer_->OnRequestTimeout((now - time_sent_).total_milliseconds());
        }
        
        // Requests must be sent no less than one second apart.
        timer_.expires_at(time_sent_ + posix_time::seconds(1));
        timer_.async_wait(boost::bind(&pinger::start_send, this));
    }
    
    void start_receive()
    {
        // Discard any data already in the buffer.
        reply_buffer_.consume(reply_buffer_.size());
        
        // Wait for a reply. We prepare the buffer to receive up to 64KB.
        socket_.async_receive(reply_buffer_.prepare(65536),
                              boost::bind(&pinger::handle_receive, this, _2));
    }
    
    void handle_receive(std::size_t length)
    {
        // The actual number of bytes received is committed to the buffer so that we
        // can extract it using a std::istream object.
        reply_buffer_.commit(length);
        
        // Decode the reply packet.
        std::istream is(&reply_buffer_);
        ipv4_header ipv4_hdr;
        icmp_header icmp_hdr;
        is >> ipv4_hdr >> icmp_hdr;
        
        // We can receive all ICMP packets received by the host, so we need to
        // filter out only the echo replies that match the our identifier and
        // expected sequence number.
        if (is && icmp_hdr.type() == icmp_header::echo_reply
            && icmp_hdr.identifier() == get_identifier()
            && icmp_hdr.sequence_number() == sequence_number_)
        {
            // If this is the first reply, interrupt the five second timeout.
            if (num_replies_++ == 0)
                timer_.cancel();
            
            // Print out some information about the reply packet.
            posix_time::ptime now = posix_time::microsec_clock::universal_time();
            
            observer_->OnReply(length - ipv4_hdr.header_length(), ipv4_hdr.time_to_live(), (now - time_sent_).total_milliseconds());
        }
        
        start_receive();
    }

    static unsigned short get_identifier()
    {
        return static_cast<unsigned short>(::getpid());
    }
    
    
    
    icmp::resolver resolver_;
    icmp::endpoint destination_;
    icmp::socket socket_;
    deadline_timer timer_;
    unsigned short sequence_number_;
    posix_time::ptime time_sent_;
    boost::asio::streambuf reply_buffer_;
    std::size_t num_replies_;
    
    unsigned int attempts_;
    unsigned int expires_at_;
    static const int milliseconds_multiplier_ = 1000;
    IPingObserver* observer_;
};

unsigned int Ping(const char* destination, unsigned int attempts, IPingObserver* observer)
{
    boost::asio::io_service io_service;
    pinger p(io_service, destination, attempts, observer);
    io_service.run();
    return p.expires_at();
}
